package com.tinyengine.it.security;

import javax.crypto.Cipher;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class RSAExample {
    
    // 生成密钥对
    public static KeyPair generateKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048); // 密钥长度
        return keyPairGenerator.generateKeyPair();
    }
    
    // 获取公钥字符串
    public static String getPublicKeyString(PublicKey publicKey) {
        return Base64.getEncoder().encodeToString(publicKey.getEncoded());
    }
    
    // 获取私钥字符串
    public static String getPrivateKeyString(PrivateKey privateKey) {
        return Base64.getEncoder().encodeToString(privateKey.getEncoded());
    }
    
    // 从字符串恢复公钥
    public static PublicKey getPublicKey(String publicKeyString) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(publicKeyString);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePublic(keySpec);
    }
    
    // 从字符串恢复私钥
    public static PrivateKey getPrivateKey(String privateKeyString) throws Exception {
        byte[] keyBytes = Base64.getDecoder().decode(privateKeyString);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return keyFactory.generatePrivate(keySpec);
    }
    
    // 使用公钥加密
    public static String encrypt(String content, String publicKeyString) throws Exception {
        PublicKey publicKey = getPublicKey(publicKeyString);
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] encryptedBytes = cipher.doFinal(content.getBytes());
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }
    
    // 使用私钥解密
    public static String decrypt(String encryptedContent, String privateKeyString) throws Exception {
        PrivateKey privateKey = getPrivateKey(privateKeyString);
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] encryptedBytes = Base64.getDecoder().decode(encryptedContent);
        byte[] decryptedBytes = cipher.doFinal(encryptedBytes);
        return new String(decryptedBytes);
    }
    
    // 测试示例
    public static void main(String[] args) {
        try {
            // 1. 生成密钥对
            KeyPair keyPair = generateKeyPair();
            String publicKeyString = getPublicKeyString(keyPair.getPublic());
            String privateKeyString = getPrivateKeyString(keyPair.getPrivate());
            
            System.out.println("公钥: " + publicKeyString);
            System.out.println("私钥: " + privateKeyString);
            
            // 2. 要加密的内容
            String content = "12342234";
            System.out.println("原始内容: " + content);
            
            // 3. 加密
            String encryptedContent = encrypt(content, publicKeyString);
            System.out.println("加密后: " + encryptedContent);
            
            // 4. 解密
            String decryptedContent = decrypt(encryptedContent, privateKeyString);
            System.out.println("解密后: " + decryptedContent);
            
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}